package org.kiwi.data.lb;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.springframework.util.Assert;

import java.util.*;

/**
 * Created by jack on 16/10/26.
 */
public class RoundRobinLoadBalanceStrategy extends AbstractLoadBalanceStrategy<String> {

    private static final int MIN_LB_FACTOR = 1;

    private int currentPos;

    private Map<String, Integer> currentTargets;
    private Map<String, Integer> failedTargets;

    public RoundRobinLoadBalanceStrategy(Map<String, Integer> lbFactors, String writeDsName) {
        Assert.notEmpty(lbFactors, "lbFactors must not be null");
        Assert.hasText(writeDsName, "writeDsName is required");

        this.writeDsName = writeDsName;
        this.currentTargets = Collections.synchronizedMap(lbFactors);
        this.failedTargets = Collections.synchronizedMap(new HashMap<String, Integer>(this.currentTargets.size()));

        reInitTargets(this.currentTargets);
    }

    private void reInitTargets(Map<String, Integer> lbFactors) {
        this.targets = initTargets(lbFactors);

        if (CollectionUtils.isEmpty(this.targets)) {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("targets is empty");
            }
        }

        this.currentPos = 0;
    }

    public List<String> initTargets(Map<String, Integer> lbFactors) {
        if (MapUtils.isEmpty(lbFactors)) {
            return null;
        }

        fixFactor(lbFactors);

        Collection<Integer> factors = lbFactors.values();

        int min = Collections.min(factors);
        if (min > MIN_LB_FACTOR && canModAll(min, factors)) {
            return buildBalanceTargets(lbFactors, min);
        }

        return buildBalanceTargets(lbFactors, MIN_LB_FACTOR);
    }

    protected synchronized List<String> getTargets() {
        if (this.targets == null) {
            this.targets = new ArrayList<>();
        }
        return this.targets;
    }

    private void fixFactor(Map<String, Integer> lbFactors) {
        Set<Map.Entry<String, Integer>> setEntries = lbFactors.entrySet();
        for (Map.Entry<String, Integer> entry : setEntries) {
            if (entry.getValue() < MIN_LB_FACTOR) {
                entry.setValue(MIN_LB_FACTOR);
            }
        }
    }

    private boolean canModAll(int base, Collection<Integer> factors) {
        for (Integer integer : factors) {
            if (integer % base != 0) {
                return false;
            }
        }
        return true;
    }

    private List<String> buildBalanceTargets(Map<String, Integer> lbFactors, int baseFactor) {
        Set<Map.Entry<String, Integer>> setEntries = lbFactors.entrySet();
        List<String> targets = new ArrayList<>();
        for (Map.Entry<String, Integer> entry : setEntries) {
            int factor = entry.getValue() / baseFactor;

            for (int i = 0; i < factor; i++) {
                targets.add(entry.getKey());
            }
        }
        return targets;
    }

    @Override
    public synchronized String elect() {
        if (CollectionUtils.isEmpty(this.targets)) {
            return null;
        }
        if (this.currentPos >= this.targets.size()) {
            this.currentPos = 0;
        }
        return this.targets.get(this.currentPos++);
    }

    @Override
    public synchronized void removeTarget(String key) {
        if (this.currentTargets.containsKey(key)) {
            this.failedTargets.put(key, this.currentTargets.get(key));
            this.currentTargets.remove(key);

            reInitTargets(this.currentTargets);
        }
    }

    @Override
    public synchronized void recoverTarget(String key) {
        if (this.failedTargets.containsKey(key)) {
            this.currentTargets.put(key, this.failedTargets.get(key));
            this.failedTargets.remove(key);

            reInitTargets(this.currentTargets);
        }
    }

    @Override
    public Set<String> getFailedTargets() {
        return this.failedTargets.keySet();
    }


}
